import torch.nn as nn
from torchvision import models

class mobilenet_v2(nn.Module):
    def __init__(self,class_number):
        super(mobilenet_v2, self).__init__()
        self.feature=models.mobilenet_v2(pretrained=False)
        # self.mobinet_layer = nn.Sequential(*list(self.feature.children())[:-1])  #detect last layer
        # self.feature.classifier =nn.Sequential( nn.Linear(in_features=1280, out_features=class_number))
        self.feature.classifier =nn.Linear(in_features=1280, out_features=class_number)

    def forward(self,x):
        y=self.feature(x)

        # y=torch.softmax(x,1)
        # y=torch.sigmoid(y)

        return y





